import bpy
import sys
import os
import argparse

def import_3d_model(model_path):
    """
    根据文件扩展名导入不同格式的3D模型
    支持: .obj, .fbx, .glb, .gltf
    """
    file_ext = os.path.splitext(model_path)[1].lower()
    
    print(f"正在导入 {file_ext} 格式的模型: {model_path}")
    
    try:
        if file_ext == '.obj':
            bpy.ops.wm.obj_import(filepath=model_path)
            print("OBJ 文件导入成功")
        elif file_ext in ['.glb', '.gltf']:
            bpy.ops.import_scene.gltf(filepath=model_path)
            print(f"{file_ext.upper()} 文件导入成功")
        else:
            raise ValueError(f"不支持的文件格式: {file_ext}")
            
    except Exception as e:
        print(f"{file_ext.upper()} 导入失败: {e}")
        raise

def setup_texture_painting(
    obj_path,              # 模型文件路径
    output_path,           # 输出文件路径
    texture_path=None,
    base_color=(1.0, 1.0, 1.0, 1.0),
    metallic=0.0,  # 金属度
    roughness=0.5  # 粗糙度
):
    # 清除当前场景中的所有对象
    bpy.ops.object.select_all(action='SELECT')
    bpy.ops.object.delete()
    
    # 检查文件是否存在
    if not os.path.exists(obj_path):
        print(f"错误: 模型文件不存在 - {obj_path}")
        return
    
    # 导入3D模型 - 支持多种格式
    try:
        import_3d_model(obj_path)
    except Exception as e:
        print(f"模型导入失败: {e}")
        return
    
    # 获取导入的对象（通常是第一个对象）
    print("正在查找导入的对象..., 一共为：", bpy.context.selected_objects)
    obj = bpy.context.selected_objects[0] if bpy.context.selected_objects else None
    
    for o in bpy.context.selected_objects:
        if o.type == 'EMPTY':
            continue
        if o.type == 'MESH':
            obj = o
            break   
        
    if not obj:
        print("错误: 没有找到有效的对象")
        return
    else:
        print(f"已找到对象: {obj.name} (类型: {obj.type})")
    
    # 选择对象
    bpy.ops.object.select_all(action='DESELECT')
    obj.select_set(True)
    bpy.context.view_layer.objects.active = obj
    
    # 检查对象是否为网格
    if obj.type != 'MESH':
        print(f"错误: 对象类型不是网格 (类型: {obj.type})")
        return
    
    # 检查纹理路径
    if texture_path:
        texture_path = os.path.normpath(texture_path)
        print(f"检查纹理路径: {texture_path}")
        
        if not os.path.exists(texture_path):
            print(f"错误: 纹理文件不存在 - {texture_path}")
            texture_path = None
    
    # 如果没有有效路径，创建临时纹理
    if not texture_path:
        print("创建临时纹理...")
        temp_dir = os.path.dirname(bpy.data.filepath) if bpy.data.filepath else os.getcwd()
        texture_path = os.path.join(temp_dir, "temp_texture.png")
        
        image = bpy.data.images.new("TempTexture", width=1024, height=1024)
        pixels = [1.0, 0.5, 0.0, 1.0] * 1024 * 1024  # 橙色
        image.pixels = pixels
        
        try:
            image.save_render(filepath=texture_path)
            print(f"临时纹理已保存到: {texture_path}")
        except Exception as e:
            print(f"无法保存临时纹理: {e}")
            texture_path = None
    
    # 创建或选择材质
    if not obj.data.materials:
        print("创建新材质")
        material = bpy.data.materials.new(name="PorcelainMaterial")
        material.use_nodes = True
        
        # 确保材质分配给对象
        obj.data.materials.append(material)
    else:
        print(f"使用现有材质: {obj.data.materials[0].name}")
        material = obj.data.materials[0]
        # 确保材质使用节点
        if not material.use_nodes:
            material.use_nodes = True
    
    # 获取材质节点
    nodes = material.node_tree.nodes
    links = material.node_tree.links
    
    # 清除默认节点（保留 Material Output）
    output_node = None
    
    for node in list(nodes):
        if node.name == "Material Output":
            output_node = node
        else:
            nodes.remove(node)
    
    # 如果没有找到输出节点，创建一个
    if not output_node:
        output_node = nodes.new('ShaderNodeOutputMaterial')
        output_node.location = (300, 0)
    
    # 创建基础节点
    bsdf_node = nodes.new('ShaderNodeBsdfPrincipled')
    bsdf_node.location = (0, 0)
    
    # 设置基础色（添加范围检查）
    for i in range(4):
        if base_color[i] > 1.0:
            print(f"警告: 颜色值 {base_color[i]} 超出范围(0.0-1.0)，将被限制")
    
    # 限制颜色值范围
    clamped_color = tuple(min(max(c, 0.0), 1.0) for c in base_color)
    bsdf_node.inputs['Base Color'].default_value = clamped_color
    print(f"设置基础色: {clamped_color}")
    
    # 设置金属度和粗糙度（帮助显示基础色）
    bsdf_node.inputs['Metallic'].default_value = metallic
    bsdf_node.inputs['Roughness'].default_value = roughness
    
    texture_node = nodes.new('ShaderNodeTexImage')
    texture_node.location = (-300, 0)
    
    # 加载纹理
    if texture_path:
        try:
            texture_node.image = bpy.data.images.load(texture_path)
            print(f"成功加载纹理: {texture_path}")
            
            # 检查纹理属性
            if texture_node.image:
                print(f"  纹理尺寸: {texture_node.image.size}")
                print(f"  纹理路径: {texture_node.image.filepath}")
        except Exception as e:
            print(f"加载纹理失败: {e}")
            texture_node.image = bpy.data.images.new("DefaultTexture", width=1024, height=1024)
    else:
        texture_node.image = bpy.data.images.new("FallbackTexture", width=1024, height=1024)
    
    # 连接节点
    try:
        links.new(texture_node.outputs['Color'], bsdf_node.inputs['Base Color'])
        links.new(bsdf_node.outputs['BSDF'], output_node.inputs['Surface'])
        print("节点连接成功")
    except Exception as e:
        print(f"节点连接失败: {e}")
        return
    
    # 确保有UV展开
    if not obj.data.uv_layers:
        print("创建UV展开...")
        bpy.ops.object.mode_set(mode='EDIT')
        bpy.ops.mesh.select_all(action='SELECT')
        bpy.ops.uv.smart_project()
        bpy.ops.object.mode_set(mode='OBJECT')
    
    # 输出UV信息
    uv_layer = obj.data.uv_layers.active
    if uv_layer:
        print(f"UV层: {uv_layer.name}")
    else:
        print("警告: 没有活动的UV层")
    
    # 切换到材质预览模式
    for area in bpy.context.screen.areas:
        if area.type == 'VIEW_3D':
            area.spaces[0].shading.type = 'MATERIAL'
            break
    
    print("=== 脚本执行完成 ===")
    
    # 创建文件夹
    if bool(os.path.dirname(output_path)):
        os.makedirs(os.path.dirname(output_path), exist_ok=True) 
    
    # 导出结果 - 支持多种输出格式
    output_ext = os.path.splitext(output_path)[1].lower()
    
    try:
        if output_ext == '.obj':
            bpy.ops.wm.obj_export(filepath=output_path)
            print(f"已保存处理后的模型为OBJ格式: {output_path}")
            
        elif output_ext in ['.glb', '.gltf']:
            # Blender 4.4.1 GLTF/GLB 导出设置
            bpy.ops.export_scene.gltf(
                filepath=output_path,
                use_selection=True
            )
            print(f"已保存处理后的模型为{output_ext.upper()}格式: {output_path}")
            
        else:
            print(f"不支持的输出格式: {output_ext}，默认保存为.blend")
            output_path = os.path.splitext(output_path)[0] + '.blend'
            bpy.ops.wm.save_as_mainfile(filepath=output_path)
            
    except Exception as e:
        print(f"导出失败: {e}")
        # fallback到blend格式
        output_path = os.path.splitext(output_path)[0] + '.blend'
        bpy.ops.wm.save_as_mainfile(filepath=output_path)
    
    print(f"最终保存路径: {output_path}")

def texture_painting_interface(obj_path, output_path, texture_path):
    """
    封装的纹理绘制接口函数
    :param obj_path: 输入模型文件路径 (支持 .obj, .glb, .gltf)
    :param output_path: 输出文件路径 (支持: .obj, .glb, .gltf)
    :param texture_path: 纹理图像路径
    """
    setup_texture_painting(
        obj_path=obj_path,
        output_path=output_path,
        texture_path=texture_path
    )

if __name__ == "__main__":
    # 解析命令行参数 - 只解析传递给Python脚本的参数
    try:
        script_args_start = sys.argv.index('--') + 1
        script_args = sys.argv[script_args_start:]
    except ValueError:
        script_args = []  # 如果没有 '--'，则没有传递给脚本的参数

    # 解析命令行参数
    parser = argparse.ArgumentParser(description='Blender纹理绘制脚本 - 支持多种3D格式')
    parser.add_argument('--obj_path', default=r"C:\Users\86184\Desktop\test-obj.obj", 
                       help='输入模型文件路径 (支持: .obj, .glb, .gltf)')
    parser.add_argument('--output_path', default=r"test_result_obj.obj", 
                       help='输出文件路径 (支持: .obj, .glb, .gltf)')
    parser.add_argument('--texture_path', default=r"C:\Users\86184\Pictures\QQ20250612-221343.jpg", 
                       help='纹理图像路径')
    
    # 将提取的脚本参数传递给解析器
    args = parser.parse_args(script_args)

    print("=== 多格式纹理绘制脚本启动 ===")
    print(f"支持的输入格式: .obj, .glb, .gltf")
    print(f"支持的输出格式: .obj, .glb, .gltf, .blend")
    print(f"输入文件: {args.obj_path}")
    print(f"输出文件: {args.output_path}")
    print(f"纹理文件: {args.texture_path}")
    print("=" * 50)

    texture_painting_interface(
        obj_path=args.obj_path,
        output_path=args.output_path,
        texture_path=args.texture_path
    )